package jm.data.structure.trie;

import jm.data.structure.hash.HashMap;

/**
 * @Description
 * @date 2022/4/30 17:20
 */
public class Trie<V> {

    private int size;

    private Node<V> root;

    /**
     * 获取元素个数
     * @return size
     */
    public int size(){
        return size;
    }

    /**
     * 判断是否为空
     * @return
     */
    public boolean isEmpty(){
        return size == 0;
    }

    /**
     * 清空集合
     */
    public void clear(){
        size = 0;
        root = null;
    }

    /**
     * 判断是否包含某个元素
     * @param key
     * @return
     */
    public boolean contains(String key){
        Node<V> node = node(key);
        return node != null && node.word;
    }

    /**
     * 根据key获取元素
     * @param key
     * @return
     */
    public V get(String key){
        Node<V> node = node(key);
        return node != null && node.word ? node.v : null;
    }

    /**
     * 添加
     * @param key
     * @param v
     * @return
     */
    public V add(String key, V v){
        keyCheck(key);
        if (root == null){
            root = new Node<>(null);
        }
        Node<V> node = root;
        int length = key.length();
        for (int i = 0; i < length; i++) {
            char c = key.charAt(i);
            // 判断 children 是否为空
            boolean emptyChildren = node.children == null;
            // 不为空就get
            Node<V> children = emptyChildren ? null : node.children.get(c);
            if (children == null){
                children = new Node<>(node);
                children.character = c;
                // 如果为空就新建一个hashMap
                node.children = emptyChildren ? new HashMap<Character,Node<V>>() : node.children;
                node.children.put(c,children);
            }
            node = children;
        }

        // 新增一个单词
        if (!node.word){
            node.word = true;
            node.v = v;
            size ++;
            return null;
        }

        V oldValue = node.v;
        node.v = v;
        return oldValue;
    }

    /**
     * 删除元素
     * @param key
     */
    public V remove(String key){
        Node<V> node = node(key);
        if (node == null || !node.word){ // 判断是否需要删除
            return null;
        }
        V oldV = node.v;
        size--;
        // 如果最后一个节点还有子节点，就无法进行删除，更改word即可
        if (node.children != null && !node.children.isEmpty()){
            node.word = false;
            node.v = null;
            return oldV;
        }

        Node<V> parent = null;
        while ((parent = node.parent) != null){
            // 从父节点中删除
            parent.children.remove(node.character);
            // 如果删除之后父节点还有其它的元素，就退出循环
            if (parent.word || !parent.children.isEmpty()){
                break;
            }
            node = parent;
        }
        return oldV;
    }

    /**
     * 判断是否存在该前驱的key
     * @param prefix
     * @return
     */
    public boolean startWith(String prefix){
        return node(prefix) != null;
    }

    /**
     * 根据key找节点
     * @param key
     * @return
     */
    private Node<V> node(String key){
        keyCheck(key);
        Node<V> node = root;
        int length = key.length();
        for (int i = 0; i < length; i++) {
            if (node == null || node.children == null || node.children.isEmpty()){
                return null;
            }
            char c = key.charAt(i);
            node = node.children.get(c);
        }
        return node;
    }

    /**
     * key 检测
     * @param key
     */
    private void keyCheck(String key){
        if (key == null || key.length() == 0){
            throw new IllegalArgumentException("key must not be empty");
        }
    }

    private static class Node<V>{
        Node<V> parent;
        HashMap<Character,Node<V>> children; // 存储子节点
        V v; // 只有在单词结尾的节点才会存储
        boolean word; // 是否为单词的结尾
        Character character;

        public Node(Node<V> parent) {
            this.parent = parent;
        }
    }
}
